package com.app.datastruct;

public class AVLTree {


    public class Node {
        public int value;
        public int height;
        public Node node_l;
        public Node node_r;

        public Node(int value, Node node_l, Node node_r) {
            this.value = value;
            this.node_l = node_l;
            this.node_r = node_r;
        }
    }

    private Node rootNode;

    public void insert(int value) {
        rootNode = insert(rootNode, value);
    }

    private Node insert(Node rootNode, int value) {
        if (rootNode == null) {
            rootNode = new Node(value, null, null);
        } else {
            if (rootNode.value > value) {   //左子树
                rootNode.node_l = insert(rootNode.node_l, value);
                if (height(rootNode.node_l) - height(rootNode.node_r) == 2) {
                    //1.LL
                    if (rootNode.node_l.value > value) {
                        rootNode = llRotation(rootNode);
                    } else {  //2.LR
                        rootNode = lrRotation(rootNode);
                    }
                }
            } else if (rootNode.value < value) {
                rootNode.node_r = insert(rootNode.node_r, value);
                if (height(rootNode.node_l) - height(rootNode.node_r) == -2) {
                    //3.RR
                    if (rootNode.node_r.value < value) {
                        rootNode = rrRotation(rootNode);
                    } else {  //4.RL
                        rootNode = rlRotation(rootNode);
                    }
                }
            }
        }
        rootNode.height = height(rootNode);
        return rootNode;
    }


    public void remove(int value) {
        remove(rootNode, value);
    }

    private Node remove(Node root, int value) {
        if (root == null) {
            return root;
        }
        if (root.value > value) {
            root.node_l = remove(root.node_l, value);
            if (height(root.node_r) - height(root.node_l) == 2) {
                //rr
                if (height(root.node_r.node_r) > 0) {
                    root = rrRotation(root);
                } else {
                    root = rlRotation(root);
                }
            }

        } else if (root.value < value) {
            root.node_r = remove(root.node_r, value);
            if (height(root.node_l) - height(root.node_r) == 2) {
                //ll
                if (height(root.node_l.node_l) > 0) {
                    root = llRotation(root);
                } else {
                    root = lrRotation(root);
                }
            }
        } else {
            if (root.node_l != null && root.node_r != null) {
                if (height(root.node_l) > height(root.node_r)) {
                    Node minNode = findMinNode(root.node_l);
                    root.value = minNode.value;
                    root.node_l = remove(root.node_l, minNode.value);
                } else {
                    Node maxNode = findMinNode(root.node_r);
                    root.value = maxNode.value;
                    root.node_r = remove(root.node_r, maxNode.value);
                }
            } else {
                root = root.node_l != null ? root.node_l : root.node_r;
            }
        }
        return root;
    }

    private Node llRotation(Node node) {
        Node tmp = null;
        tmp = node.node_l;
        node.node_l = tmp.node_r;
        tmp.node_r = node;

        tmp.height = height(tmp);
        node.height = height(node);
        return tmp;
    }

    private Node rrRotation(Node node) {
        Node tmp = null;
        tmp = node.node_r;
        node.node_r = tmp.node_l;
        tmp.node_l = node;

        tmp.height = height(tmp);
        node.height = height(node);
        return tmp;

    }

    private Node lrRotation(Node node) {
        node.node_l = rrRotation(node.node_l);
        return llRotation(node);
    }

    private Node rlRotation(Node node) {
        node.node_r = llRotation(node.node_r);
        return rrRotation(node);
    }


    private int height(Node node) {
        if (node == null) {
            return 0;
        } else {
            return Math.max(height(node.node_l), height(node.node_r)) + 1;
        }
    }


    public Node search(int value) {
        return search(rootNode, value);
    }

    private Node findMinNode(Node rootNode) {
        Node minNode = null;
        while (rootNode != null) {
            minNode = rootNode;
            rootNode = rootNode.node_l;
        }
        return minNode;
    }

    private Node findMaxNode(Node rootNode) {
        Node maxNode = null;
        while (rootNode != null) {
            maxNode = rootNode;
            rootNode = rootNode.node_r;
        }
        return maxNode;
    }

    private Node search(Node root, int value) {
        if (root == null) {
            return null;
        } else {
            if (root.value > value) {
                return search(root.node_l, value);
            } else if (root.value < value) {
                return search(root.node_r, value);
            } else {
                return root;
            }
        }
    }

    public static void main(String[] args) {
        AVLTree tree = new AVLTree();
        tree.insert(10);
        tree.insert(9);
        tree.insert(8);
        tree.insert(7);
        tree.insert(6);
        tree.insert(5);
        tree.insert(4);
        tree.insert(20);
        tree.remove(8);


        Node node = tree.rootNode;
        while (node != null) {
            System.out.println(node.value);
            node = node.node_r;
        }


    }
//    https://www.cnblogs.com/skywang12345/p/3577479.html#a2
}
